# flask_server.py
from flask import Flask, request, jsonify, json
from resemblyzer import VoiceEncoder, preprocess_wav
import os
import uuid
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from datasets import load_dataset

app = Flask(__name__)

# 检查是否有 GPU，如果有则使用 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device} - {torch.cuda.get_device_name(0)}")
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# 模型链接：https://www.modelscope.cn/models/openai-mirror/whisper-large-v3-turbo
model_id = "openai/whisper-large-v3-turbo"

# 初始化 VoiceEncoder 时指定设备
encoder = VoiceEncoder(device=device)

# 声纹数据存储路径
voice_db_path = "voice_db"
os.makedirs(voice_db_path, exist_ok=True)

@app.route("/enroll", methods=["POST"])
def enroll():
    """
    注册声纹的函数。

    该函数接收一个音频文件和一个用户名，生成一个唯一的ID，然后将音频文件转换为声纹特征，
    最后将用户的ID、用户名和声纹特征保存到JSON文件中。

    返回值:
        jsonify: 包含用户ID和用户名的JSON响应。
    """
    # 从请求中获取音频文件
    file = request.files["audio"]
    # 从请求中获取用户名
    name = request.form["name"]

    # 生成一个唯一的文件名，并将音频文件保存到指定路径
    audio_path = os.path.join(voice_db_path, f"{uuid.uuid4()}.wav")
    file.save(audio_path)

    # 加载音频文件并生成声纹特征
    # embedding = encoder.embed_utterance(encoder.load_wav(audio_path))
    # 使用 preprocess_wav 加载音频数据
    wav = preprocess_wav(audio_path)

    # 生成声纹嵌入
    embedding = encoder.embed_utterance(wav)
    
    # 删除临时音频文件
    os.remove(audio_path)

    # 生成一个唯一的用户ID
    user_id = str(uuid.uuid4())
    # 创建一个包含用户ID、用户名和声纹特征的字典
    user_data = {"id": user_id, "name": name, "embedding": embedding.tolist()}

    # 将用户数据保存到JSON文件中
    with open(os.path.join(voice_db_path, f"{user_id}.json"), "w") as f:
        import json
        json.dump(user_data, f)

    # 返回包含用户ID和用户名的JSON响应
    return jsonify({"id": user_id, "name": name})


@app.route("/identify", methods=["POST"])
def identify():
    """
    识别声纹的函数。

    该函数接收一个音频文件，生成一个唯一的ID，然后将音频文件转换为声纹特征，
    最后将用户的ID、用户名和声纹特征保存到JSON文件中。

    返回值:
        jsonify: 包含用户ID和用户名的JSON响应。
    """
    # 从请求中获取音频文件
    file = request.files["audio"]
    # 生成一个唯一的文件名，并将音频文件保存到指定路径
    audio_path = os.path.join(voice_db_path, f"{uuid.uuid4()}.wav")
    file.save(audio_path)

    # 加载音频文件并生成声纹特征
    # test_embedding = encoder.embed_utterance(encoder.load_wav(audio_path))
    # 使用 preprocess_wav 加载音频数据
    wav = preprocess_wav(audio_path)
    # 生成声纹嵌入
    test_embedding = encoder.embed_utterance(wav)

    # 删除临时音频文件
    os.remove(audio_path)

    # 加载所有声纹并进行比较
    from scipy.spatial.distance import cosine
    min_distance = float("inf")
    identified_user = None

    # 遍历声纹数据库中的所有用户
    for file in os.listdir(voice_db_path):
        # 只处理JSON文件
        if file.endswith(".json"):
            # 打开JSON文件并加载用户数据
            with open(os.path.join(voice_db_path, file), "r") as f:
                user_data = json.load(f)
                # 计算测试声纹与数据库中声纹的余弦距离
                dist = cosine(user_data["embedding"], test_embedding)
                # 如果距离小于当前最小距离，则更新最小距离和识别的用户
                if dist < min_distance:
                    min_distance = dist
                    identified_user = user_data

    # 如果识别到用户并且距离小于阈值，则返回用户ID和用户名
    if identified_user and min_distance < 0.5:  # 距离阈值
        return jsonify({"id": identified_user["id"], "name": identified_user["name"]})
    # 否则返回错误信息
    else:
        return jsonify({"error": "User not identified"}), 404
    
def get_all_users():
    """
    获取所有用户数据的函数。

    该函数遍历声纹数据库目录中的所有JSON文件，加载每个文件中的用户数据，
    并将其添加到一个列表中。最后返回包含所有用户数据的列表。

    返回值:
        list: 包含所有用户数据的列表。
    """
    users = []
    # 遍历声纹数据库目录中的所有文件
    for file in os.listdir(voice_db_path):
        # 只处理JSON文件
        if file.endswith(".json"):
            # 打开JSON文件并加载用户数据
            with open(os.path.join(voice_db_path, file), "r") as f:
                user_data = json.load(f)
                # 将用户数据添加到列表中
                users.append(user_data)
    # 返回包含所有用户数据的列表
    return users


# 获取所有用户
@app.route('/get_users', methods=['GET'])
def get_users():
    """
    获取分页用户数据的函数。

    该函数首先调用 get_all_users 函数获取所有用户数据，然后根据请求参数中的页码和每页显示数量，
    计算出当前页的用户数据，并构建一个包含页码、总页数和当前页用户数据的 JSON 响应。

    返回值:
        jsonify: 包含页码、总页数和当前页用户数据的 JSON 响应。
    """
    # 调用函数获取所有用户列表
    all_users = get_all_users()

    # 获取页数参数，默认第一页
    page = int(request.args.get('page', 1))
    # 每页显示数量
    per_page = 5
    # 计算总页面数
    total_pages = (len(all_users) + per_page - 1) // per_page
    # 计算起始索引
    start_index = (page - 1) * per_page
    # 获取当前页的用户数据
    current_page_users = all_users[start_index:start_index + per_page]
    # 构建返回结果
    result = {
        'page': page,
        'total_pages': total_pages,
        'users': [{'id': user['id'], 'name': user['name']} for user in current_page_users]
    }
    return jsonify(result)


# 删除用户
@app.route('/delete_user', methods=['DELETE'])
def delete_user():
    """
    删除用户的函数。

    该函数接收一个用户ID作为参数，然后检查该用户是否存在。如果用户存在，则删除对应的JSON文件；
    如果用户不存在，则返回错误信息。

    返回值:
        jsonify: 包含操作结果的JSON响应。
    """
    # 获取用户ID参数
    user_id = request.args.get('user_id')
    
    # 检查用户ID是否存在
    if not user_id:
        return jsonify({"error": "Missing user_id parameter"}), 400
    
    # 构建用户JSON文件路径
    user_file_path = os.path.join(voice_db_path, f"{user_id}.json")
    
    # 检查文件是否存在
    if not os.path.exists(user_file_path):
        return jsonify({"error": "User not found"}), 404
    
    # 删除用户JSON文件
    os.remove(user_file_path)
    
    return jsonify({"message": "User deleted successfully"}), 200

@app.route("/transcribe", methods=["POST"])
def transcribe():
    """
    转录音频文件的函数。

    该函数接收一个音频文件，生成一个唯一的ID，然后将音频文件转换为文本。
    最后返回包含转录文本的JSON响应。

    返回值:
        jsonify: 包含转录文本的JSON响应。
    """
    # 从请求中获取音频文件
    file = request.files["audio"]
    # 生成一个唯一的文件名，并将音频文件保存到指定路径
    audio_path = os.path.join(voice_db_path, f"{uuid.uuid4()}.wav")
    file.save(audio_path)

    # 使用 Whisper 模型转录音频文件
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
    )
    model.to(device)

    processor = AutoProcessor.from_pretrained(model_id)

    pipe = pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        torch_dtype=torch_dtype,
        device=device,
    )

    result = pipe(audio_path)
    
    transcription = result["text"]

    # 删除临时音频文件
    os.remove(audio_path)

    # 返回包含转录文本的JSON响应
    return jsonify({"transcription": transcription})

if __name__ == "__main__":
    app.run(debug=True, port=8080)